package unionset

/*
	并查集，需要拿到原数据，包装成节点进行初始化
		1. 方法：
			Union(a, b) 	将两个节点连通 O(1)
			IsSameSet(a, b) 判断两个节点是否是在同一个连通区 O(1)
			SetNum()  		获取到未连通区域的数量O(1)
			MaxSetSize()	获取最大的连通区间的节点数

*/

type Node struct {
	Value interface{} // 存的实际节点的值，用于包装
}

type UnionSet struct {
	nodeMap   map[interface{}]Node // 保存值与节点的映射关系
	parentMap map[Node]Node        // 保存节点的父子映射关系
	sizeMap   map[Node]int		   // 保存集合的大小数据，与连通区的多少
	maxUSize  int                  // 保存集合的大小数据，与连通区的多少
}

func NewUnionSet(datas []interface{}) UnionSet {
	us := UnionSet{}
	us.maxUSize = 1
	us.parentMap = make(map[Node]Node)
	us.nodeMap = make(map[interface{}]Node)
	us.sizeMap = make(map[Node]int)
	for i:=0; i<len(datas); i++ {
		us.nodeMap[datas[i]] = Node{datas[i]}
	}
	for i:=0; i<len(us.nodeMap); i++ {
		us.parentMap[us.nodeMap[datas[i]]] = us.nodeMap[datas[i]]
		us.sizeMap[us.nodeMap[datas[i]]] = 1
	}
	return us
}


func (us *UnionSet) Union(a, b interface{}){
	if !us.Contain(a) || !us.Contain(b) {
		panic(any("not contain value input"))
	}
	if !us.IsSameSet(a, b){
		parentA := us.FindParent(us.nodeMap[a])
		parentB := us.FindParent(us.nodeMap[b])
		if us.sizeMap[parentA] > us.sizeMap[parentB] {
			us.parentMap[parentB] = parentA
			us.sizeMap[parentA] += us.sizeMap[parentB]
			us.maxUSize = max(us.maxUSize, us.sizeMap[parentA])
			delete(us.sizeMap, parentB)
		} else {
			us.parentMap[parentA] = parentB
			us.sizeMap[parentB] += us.sizeMap[parentA]
			us.maxUSize = max(us.maxUSize, us.sizeMap[parentB])
			delete(us.sizeMap, parentA)
		}
	}
}

func (us *UnionSet) IsSameSet(a, b interface{}) bool {
	if !us.Contain(a) || !us.Contain(b) {
		panic(any("not contain value input"))
	}
	parentA := us.FindParent(us.nodeMap[a])
	parentB := us.FindParent(us.nodeMap[b])
	return parentA == parentB
}

func (us *UnionSet) FindParent(node Node) Node {
	nodePath := make([]Node, 0)
	for node != us.parentMap[node] {
		nodePath = append(nodePath, node)
		node = us.parentMap[node]
	}
	for i:=0; i<len(nodePath); i++ {
		us.parentMap[nodePath[i]] = node
	}
	return node
}

func (us *UnionSet) Contain(v interface{}) bool {
	if _, ok := us.nodeMap[v];ok {
		return true
	}
	return false
}

func (us *UnionSet) UnionSize() int {
	return len(us.sizeMap)
}

func (us *UnionSet) MaxUnionSize() int {
	return us.maxUSize
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}